"""HuggingFace LLM实现（适配DeepSeek）"""
from typing import Optional, List
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
import torch
from src.llm.base_llm import BaseLLM

class HuggingFaceLLM(BaseLLM):
    def __init__(
        self,
        model_name: str,
        model_path: str,
        device: str = "auto",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
        quantization: int = 4
    ):
        super().__init__()
        self.model_name = model_name
        self.model_path = model_path
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.device = device if device != "auto" else ("cuda" if torch.cuda.is_available() else "cpu")

        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path if self._is_local_path(model_path) else model_name,
            trust_remote_code=True
        )

        # 配置量化
        quantization_config = None
        if quantization in [4, 8]:
            quantization_config = BitsAndBytesConfig(
                load_in_4bit=quantization == 4,
                load_in_8bit=quantization == 8,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4"
            )

        # 加载模型
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path if self._is_local_path(model_path) else model_name,
            quantization_config=quantization_config,
            device_map="auto" if self.device == "cuda" else self.device,
            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
            trust_remote_code=True
        )

        # 构建生成管道
        self.pipeline = pipeline(
            "text-generation",
            model=self.model,
            tokenizer=self.tokenizer,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            pad_token_id=self.tokenizer.eos_token_id
        )

    def _is_local_path(self, path: str) -> bool:
        return path.startswith("./") or path.startswith("/") or ":" in path

    def generate(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        # DeepSeek对话格式
        if "deepseek" in self.model_name.lower() or "chat" in self.model_name.lower():
            prompt = f"### User: {prompt}\n### Assistant: "

        result = self.pipeline(prompt, stop=stop or ["### User:", "\n\n"])
        generated_text = result[0]["generated_text"][len(prompt):].strip()
        return generated_text